import math
import os
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from matplotlib.lines import Line2D
from torch import Tensor


def create_interactive_masks(initial_observation: dict, gui_save_dir: str = "gui_interactions") -> dict:
    """
    使用 Matplotlib 在主进程中创建交互式GUI，让用户绘制忽略区域的掩码。

    Args:
        initial_observation (dict): 包含图像数据的观测字典。
                                     图像键应以 "observation.images." 开头。
        gui_save_dir (str): 保存GUI交互快照的目录。

    Returns:
        dict: 一个字典，键是相机名，值是用户绘制的像素级掩码 (NumPy数组)。
              如果用户放弃操作或未绘制任何内容，则返回空字典。
    """
    print("\n--- 启动 Matplotlib GUI 进行交互式掩码创建... ---")

    cam_images = {}
    # 获取所有图像键
    cam_keys = [key for key in initial_observation.keys() if key.startswith("observation.images.")]

    if not cam_keys:
        print("--- 警告：在观测数据中未找到图像，跳过交互。---")
        return {}

    # 读取图像数据
    for cam_key in cam_keys:
        img_tensor = initial_observation[cam_key][0].cpu()
        img_np_rgb = (img_tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
        cam_images[cam_key] = img_np_rgb

    # --- 内部类：绘图逻辑 ---
    class MaskDrawer:
        def __init__(self, ax):
            self.ax = ax
            self.mask = None
            self.press = None
            self.lines = []

        def set_image(self, image):
            self.mask = np.zeros(image.shape[:2], dtype=np.uint8)

        def connect(self):
            self.cid_press = self.ax.figure.canvas.mpl_connect('button_press_event', self.on_press)
            self.cid_release = self.ax.figure.canvas.mpl_connect('button_release_event', self.on_release)
            self.cid_motion = self.ax.figure.canvas.mpl_connect('motion_notify_event', self.on_motion)

        def on_press(self, event):
            if event.inaxes != self.ax: return
            self.press = (event.xdata, event.ydata)

        def on_motion(self, event):
            if self.press is None or event.inaxes != self.ax: return
            x, y = event.xdata, event.ydata
            # 绘制视觉线段
            line = Line2D([self.press[0], x], [self.press[1], y], color='white', linewidth=10, alpha=0.6)
            self.ax.add_line(line)
            self.lines.append(line)
            # 在掩码上绘制
            cv2.line(self.mask, (int(self.press[0]), int(self.press[1])), (int(x), int(y)), 255, thickness=15)
            self.press = (x, y)
            self.ax.figure.canvas.draw()

        def on_release(self, event):
            self.press = None

        def clear_drawing(self):
            for line in self.lines:
                line.remove()
            self.lines.clear()
            if self.mask is not None:
                self.mask[:] = 0
            self.ax.figure.canvas.draw()

    # 在创建绘图前，全局禁用工具栏
    plt.rcParams['toolbar'] = 'none'

    num_cameras = len(cam_images)

    # 如果摄像头数量大于4，使用3列，否则使用2列
    max_cols = 3 if num_cameras > 4 else 2

    # 核心修复：使用向上取整计算行数
    num_rows = math.ceil(num_cameras / max_cols)

    # 动态调整 figsize，防止图像太小
    fig_width = 5 * max_cols
    fig_height = 4.5 * num_rows
    fig, axes = plt.subplots(num_rows, max_cols, figsize=(fig_width, fig_height), squeeze=False)

    # 尝试最大化窗口
    try:
        manager = plt.get_current_fig_manager()
        if hasattr(manager, 'window') and hasattr(manager.window, 'showMaximized'):
            manager.window.showMaximized()  # Qt backend
        elif hasattr(manager, 'resize'):
            manager.resize(*manager.window.maxsize())  # Tk backend
    except:
        pass

    axes_flat = axes.flatten()

    fig.suptitle("Interactive Mask Creator\nDraw: Drag Mouse | Save: 's' or Enter | Reset: 'r' | Quit: 'q'",
                 fontsize=14)

    drawers = {}
    cam_items = list(cam_images.items())

    # 遍历所有格子，处理多余的空白格
    for i, ax in enumerate(axes_flat):
        if i < num_cameras:
            # 这是一个有效的摄像头位置
            cam_key, img_np_rgb = cam_items[i]

            ax.imshow(img_np_rgb)
            ax.set_title(cam_key.replace("observation.images.", ""), fontsize=10)  # 简化标题
            ax.set_xticks([])
            ax.set_yticks([])

            drawer = MaskDrawer(ax)
            drawer.set_image(img_np_rgb)
            drawer.connect()
            drawers[cam_key] = drawer
        else:
            # 这是一个多余的格子（例如3个摄像头，2x2布局，第4个格子）
            ax.axis('off')  # 隐藏坐标轴
            ax.set_visible(False)  # 隐藏整个子图

    final_user_masks = {}

    def on_key_press(event):
        nonlocal final_user_masks
        if event.key == 's' or event.key == 'enter':
            os.makedirs(gui_save_dir, exist_ok=True)
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            save_path = os.path.join(gui_save_dir, f"interaction_{timestamp}.png")
            print(f"正在将GUI交互图像保存到: {save_path}")
            try:
                fig.savefig(save_path, bbox_inches='tight', dpi=150)
                print("图像保存成功。")
            except Exception as e:
                print(f"错误：无法保存GUI图像: {e}")

            print("掩码已保存，关闭窗口...")
            final_user_masks = {key: drawer.mask for key, drawer in drawers.items() if np.sum(drawer.mask) > 0}
            plt.close(fig)

        elif event.key == 'q':
            print("操作已放弃，关闭窗口...")
            final_user_masks = {}
            plt.close(fig)
        elif event.key == 'r':
            print("所有掩码已被重置，请重新绘制。")
            for drawer in drawers.values():
                drawer.clear_drawing()

    fig.canvas.mpl_connect('key_press_event', on_key_press)

    # 调整布局，避免标题重叠
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.show()

    print("--- GUI 交互已结束 ---")
    return final_user_masks


@torch.no_grad()
def process_pixel_mask_to_feature_mask(
        mask_image: np.ndarray,
        feature_map_size: tuple[int, int],
) -> Tensor | None:
    """
    将像素级掩码转换为扁平化的、适用于注意力机制的特征级掩码。
    """
    if mask_image is None or feature_map_size is None:
        return None

    h_feat, w_feat = feature_map_size

    # 将灰度掩码缩放到特征图尺寸
    mask_resized = cv2.resize(mask_image, (w_feat, h_feat), interpolation=cv2.INTER_NEAREST)

    # 转换为布尔张量 (True 表示被遮挡) 并展平
    mask_tensor = torch.from_numpy(mask_resized > 0).flatten()

    print(f"处理了注意力掩码。共 {mask_tensor.sum()} 个视觉token将被忽略。")
    return mask_tensor
